#include "ego_planner_adapter/target_relay.hpp"
#include <tf2/utils.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>

using namespace std::chrono_literals;

namespace ego_planner_adapter {

TargetRelay::TargetRelay(const rclcpp::NodeOptions& options)
  : Node("target_relay", options),
    target_altitude_(1.0),
    has_altitude_(false),
    has_pending_goal_(false),
    has_current_goal_(false),
    marker_id_(0) {
    
    RCLCPP_INFO(get_logger(), "TargetRelay node created");
    
    // Declare parameters
    declare_parameter("map_frame", "map");
    declare_parameter("base_frame", "base_footprint");
    declare_parameter("default_altitude", 1.0);
    declare_parameter("goal_tolerance", 0.5);
    declare_parameter("auto_trigger", true);
    declare_parameter("marker_scale", 0.2);
    
    // Get parameters
    map_frame_ = get_parameter("map_frame").as_string();
    base_frame_ = get_parameter("base_frame").as_string();
    default_altitude_ = get_parameter("default_altitude").as_double();
    goal_tolerance_ = get_parameter("goal_tolerance").as_double();
    auto_trigger_ = get_parameter("auto_trigger").as_bool();
    marker_scale_ = get_parameter("marker_scale").as_double();
    
    // Initialize TF2
    tf_buffer_ = std::make_shared<tf2_ros::Buffer>(get_clock());
    tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);
    
    // Create subscribers
    rviz_goal_sub_ = create_subscription<geometry_msgs::msg::PoseStamped>(
        "goal_pose", 10,
        std::bind(&TargetRelay::rvizGoalCallback, this, std::placeholders::_1)
    );
    
    target_altitude_sub_ = create_subscription<std_msgs::msg::Float64>(
        "target_altitude", 10,
        std::bind(&TargetRelay::targetAltitudeCallback, this, std::placeholders::_1)
    );
    
    trigger_sub_ = create_subscription<std_msgs::msg::Bool>(
        "planning_trigger", 10,
        std::bind(&TargetRelay::triggerCallback, this, std::placeholders::_1)
    );
    
    // Create publishers
    ego_goal_pub_ = create_publisher<geometry_msgs::msg::PoseStamped>(
        "waypoint_generator/waypoints", 10
    );
    
    goal_viz_pub_ = create_publisher<visualization_msgs::msg::Marker>(
        "goal_marker", 10
    );
    
    planning_trigger_pub_ = create_publisher<std_msgs::msg::Bool>(
        "planning_trigger", 10
    );
    
    // Initialize with default altitude
    target_altitude_ = default_altitude_;
    has_altitude_ = true;
    
    RCLCPP_INFO(get_logger(), "TargetRelay initialized with default altitude: %.2f m", default_altitude_);
}

void TargetRelay::rvizGoalCallback(const std::shared_ptr<const geometry_msgs::msg::PoseStamped>& msg) {
    RCLCPP_INFO(get_logger(), "Received 2D goal from RViz at [%.2f, %.2f] in frame '%s'",
               msg->pose.position.x, msg->pose.position.y, msg->header.frame_id.c_str());
    
    // Store pending goal
    pending_goal_ = *msg;
    has_pending_goal_ = true;
    last_goal_time_ = now();
    
    // Process immediately if we have altitude info
    if (has_altitude_) {
        processGoal(pending_goal_);
        has_pending_goal_ = false;
    } else {
        RCLCPP_WARN(get_logger(), "No altitude information available, using default: %.2f m", default_altitude_);
        target_altitude_ = default_altitude_;
        has_altitude_ = true;
        processGoal(pending_goal_);
        has_pending_goal_ = false;
    }
}

void TargetRelay::targetAltitudeCallback(const std::shared_ptr<const std_msgs::msg::Float64>& msg) {
    double new_altitude = msg->data;
    
    // Validate altitude
    if (new_altitude < 0.1 || new_altitude > 10.0) {
        RCLCPP_WARN(get_logger(), "Invalid altitude %.2f m, using previous value %.2f m", 
                   new_altitude, target_altitude_);
        return;
    }
    
    target_altitude_ = new_altitude;
    has_altitude_ = true;
    
    RCLCPP_INFO(get_logger(), "Updated target altitude to %.2f m", target_altitude_);
    
    // Process pending goal if available
    if (has_pending_goal_) {
        processGoal(pending_goal_);
        has_pending_goal_ = false;
    }
}

void TargetRelay::triggerCallback(const std::shared_ptr<const std_msgs::msg::Bool>& msg) {
    if (msg->data && has_current_goal_) {
        RCLCPP_INFO(get_logger(), "Planning trigger received, republishing current goal");
        
        // Republish current goal
        auto goal_msg = std::make_unique<geometry_msgs::msg::PoseStamped>(current_goal_);
        goal_msg->header.stamp = now();
        ego_goal_pub_->publish(std::move(goal_msg));
        
        // Refresh visualization
        publishGoalVisualization(current_goal_);
    } else if (msg->data) {
        RCLCPP_WARN(get_logger(), "Planning trigger received but no current goal available");
    }
}

void TargetRelay::processGoal(const geometry_msgs::msg::PoseStamped& goal_2d) {
    // Create 3D goal with altitude
    geometry_msgs::msg::PoseStamped goal_3d = goal_2d;
    goal_3d.pose.position.z = target_altitude_;
    goal_3d.header.stamp = now();
    
    // Transform to map frame if necessary
    geometry_msgs::msg::PoseStamped goal_map;
    if (!transformPose(goal_3d, goal_map, map_frame_)) {
        RCLCPP_ERROR(get_logger(), "Failed to transform goal to map frame");
        return;
    }
    
    // Check if goal is significantly different from current goal
    if (has_current_goal_) {
        double dx = goal_map.pose.position.x - current_goal_.pose.position.x;
        double dy = goal_map.pose.position.y - current_goal_.pose.position.y;
        double dz = goal_map.pose.position.z - current_goal_.pose.position.z;
        double distance = std::sqrt(dx*dx + dy*dy + dz*dz);
        
        if (distance < MIN_GOAL_DISTANCE) {
            RCLCPP_INFO(get_logger(), "Goal too close to current goal (%.3f m), ignoring", distance);
            return;
        }
    }
    
    // Store as current goal
    current_goal_ = goal_map;
    has_current_goal_ = true;
    
    RCLCPP_INFO(get_logger(), "Processing 3D goal: [%.2f, %.2f, %.2f] in frame '%s'",
               goal_map.pose.position.x, goal_map.pose.position.y, goal_map.pose.position.z,
               goal_map.header.frame_id.c_str());
    
    // Publish to ego-planner
    auto ego_goal_msg = std::make_unique<geometry_msgs::msg::PoseStamped>(goal_map);
    ego_goal_pub_->publish(std::move(ego_goal_msg));
    
    // Publish visualization
    publishGoalVisualization(goal_map);
    
    // Auto-trigger planning if enabled
    if (auto_trigger_) {
        auto trigger_msg = std::make_unique<std_msgs::msg::Bool>();
        trigger_msg->data = true;
        planning_trigger_pub_->publish(std::move(trigger_msg));
        
        RCLCPP_INFO(get_logger(), "Auto-triggered ego-planner");
    }
}

void TargetRelay::publishGoalVisualization(const geometry_msgs::msg::PoseStamped& goal_3d) {
    auto marker = std::make_unique<visualization_msgs::msg::Marker>();
    
    marker->header = goal_3d.header;
    marker->ns = "ego_planner_goals";
    marker->id = marker_id_++;
    marker->type = visualization_msgs::msg::Marker::ARROW;
    marker->action = visualization_msgs::msg::Marker::ADD;
    
    // Set pose
    marker->pose = goal_3d.pose;
    
    // Set scale
    marker->scale.x = marker_scale_;
    marker->scale.y = marker_scale_ * 0.5;
    marker->scale.z = marker_scale_ * 0.5;
    
    // Set color (green for ego-planner goals)
    marker->color.r = 0.0;
    marker->color.g = 1.0;
    marker->color.b = 0.0;
    marker->color.a = 1.0;
    
    // Set lifetime
    marker->lifetime = rclcpp::Duration::from_seconds(GOAL_TIMEOUT);
    
    goal_viz_pub_->publish(std::move(marker));
    
    RCLCPP_DEBUG(get_logger(), "Published goal visualization marker with ID %d", marker_id_ - 1);
}

bool TargetRelay::transformPose(const geometry_msgs::msg::PoseStamped& input_pose,
                               geometry_msgs::msg::PoseStamped& output_pose,
                               const std::string& target_frame) {
    
    if (input_pose.header.frame_id == target_frame) {
        output_pose = input_pose;
        return true;
    }
    
    try {
        // Check if transform is available
        if (!tf_buffer_->canTransform(target_frame, input_pose.header.frame_id, 
                                     input_pose.header.stamp, rclcpp::Duration::from_seconds(1.0))) {
            RCLCPP_WARN(get_logger(), "Transform from '%s' to '%s' not available",
                       input_pose.header.frame_id.c_str(), target_frame.c_str());
            return false;
        }
        
        // Perform transformation
        tf_buffer_->transform(input_pose, output_pose, target_frame);
        
        RCLCPP_DEBUG(get_logger(), "Transformed pose from '%s' to '%s'",
                    input_pose.header.frame_id.c_str(), target_frame.c_str());
        
        return true;
        
    } catch (const tf2::TransformException& ex) {
        RCLCPP_ERROR(get_logger(), "Transform failed: %s", ex.what());
        return false;
    }
}

} // namespace ego_planner_adapter